#include "include/bev_fusion.h"

#include <stdio.h>

#include <string>

#include "utils.h"

BevFusion::BevFusion() noexcept {
  // grid_buffer_ = std::make_unique<std::array<char, kGridSampleResSize>>();
  grid_buffer_ = new char[kGridSampleResSize];
  // image_out_buf_.cat2491 = std::make_shared<std::array<char, kCat2491>>();
  // image_out_buf_.cat2492 = std::make_shared<std::array<char, kCat2492>>();
  image_out_buf_.cat2491 = new char[kCat2491];
  image_out_buf_.cat2492 = new char[kCat2492];
  Print(".image_out_buf_.cat2491 = %p\n", image_out_buf_.cat2491);
}
BevFusion::~BevFusion() noexcept {
  // if (image_out_buf_.cat2492 != nullptr) {
  //   free(image_out_buf_.cat2491);
  //   image_out_buf_.cat2492 = nullptr;
  // }
  delete image_out_buf_.cat2491;
  delete image_out_buf_.cat2492;
  delete grid_buffer_;
}

uint32_t BevFusion::Init(Config const &config) noexcept {
  // step1 初始化
  Print("Initializing bevfusion ...\n");
  // 1. ImageEncoder
  uint32_t Rval = image_encoder.Init(config.image_encoder_path);
  if (Rval != 0U) {
    return 1;
  }
  Print("image_encoder init success ...\n");

  // 2. ImageEncoder
  Rval = grid_sample.Init(config.grid_sample_path);
  if (Rval != 0U) {
    return 1;
  }
  Print("grid_sample init success ...\n");

  Rval = bev_net.Init(config.bev_net_path);
  if (Rval != 0U) {
    Print("Error encounted when loading and opening flexidag bin! Quit...\n");
    return 1;
  }
  Print("Initializing bevfusion success end...\n");

  std::uint32_t box_reg_len =
      config.out_location.box_loc[0][InferBoxType::REG].prev_dim *
      config.out_location.box_loc[0][InferBoxType::REG].last_dim;
  std::uint32_t box_cls_len =
      config.out_location.box_loc[0][InferBoxType::CLS].prev_dim *
      config.out_location.box_loc[0][InferBoxType::CLS].last_dim;
  std::uint32_t box_rot_len =
      config.out_location.box_loc[0][InferBoxType::ROT].prev_dim *
      config.out_location.box_loc[0][InferBoxType::ROT].last_dim;
  std::uint32_t box_dim_len =
      config.out_location.box_loc[0][InferBoxType::DIM].prev_dim *
      config.out_location.box_loc[0][InferBoxType::DIM].last_dim;
  std::uint32_t box_vel_len =
      config.out_location.box_loc[0][InferBoxType::VEL].prev_dim *
      config.out_location.box_loc[0][InferBoxType::VEL].last_dim;
  std::uint32_t box_height_len =
      config.out_location.box_loc[0][InferBoxType::HEIGHT].prev_dim *
      config.out_location.box_loc[0][InferBoxType::HEIGHT].last_dim;

  for (std::uint8_t i = 0; i < 6; i++) {
    output_data_raw_.box[i].reg.reserve(box_reg_len);
    output_data_raw_.box[i].cls.reserve(box_cls_len);
    output_data_raw_.box[i].dim.reserve(box_dim_len);
    output_data_raw_.box[i].rot.reserve(box_rot_len);
    output_data_raw_.box[i].vel.reserve(box_vel_len);
    output_data_raw_.box[i].height.reserve(box_height_len);
  }

  std::uint32_t map_reg_len =
      config.out_location.map_loc[0][InferBoxType::REG].prev_dim *
      config.out_location.map_loc[0][InferBoxType::REG].last_dim;
  std::uint32_t map_cls_len =
      config.out_location.map_loc[0][InferBoxType::CLS].prev_dim *
      config.out_location.map_loc[0][InferBoxType::CLS].last_dim;

  for (std::uint8_t i = 0; i < 3; i++) {
    output_data_raw_.map[i].reg.reserve(map_reg_len);
    output_data_raw_.map[i].cls.reserve(map_cls_len);   
  }
  
  return Rval;
}

uint32_t BevFusion::Run(DataBuf &data_frame) noexcept {
  for (std::uint8_t i = 0; i < 6; ++i) {
    bool flag = image_encoder.SetImageFromBuf(
        data_frame.camera_data[0].first, data_frame.camera_data[0].second, i);
    if (!flag) {
      Print("Error encounted when SetImageFromBuf! Quit...\n");
      return 1;
    }
  }

  Print("SetImageFromBuf success ...\n");
  image_encoder.alloc_IO_buffer();

  bool success{true};
  success = grid_sample.SetInput(data_frame.pre_input.first,
                                 data_frame.pre_input.second,
                                 GridSample::InputType::PRE_INPUT);
  if (!success) {
    Print("Error encounted when grid_sample.SetInput pre_input! Quit...\n");
    return 1;
  }
  success = grid_sample.SetInput(data_frame.align_matrix.first,
                                 data_frame.align_matrix.second,
                                 GridSample::InputType::ALIGN_MATRIX);
  if (!success) {
    Print("Error encounted when grid_sample.SetInput align_matrix! Quit...\n");
    return 1;
  }
  grid_sample.alloc_IO_buffer();
  Print("grid_sample.SetInput success ...\n");

  success = bev_net.SetInput(data_frame.ring_matrix.first,
                             data_frame.ring_matrix.second,
                             BevNet::InputType::RING_MATRIX);
  success = bev_net.SetInput(data_frame.ray_matrix.first,
                             data_frame.ray_matrix.second,
                             BevNet::InputType::RAY_MATRIX);
  success = bev_net.SetInput(data_frame.lidar_data.first,
                             data_frame.lidar_data.second,
                             BevNet::InputType::LIDAR_INPUT);

  if (!success) {
    Print("Error encounted when load input file bev net! Quit...\n");
    return 1;
  }
  Print("bev_net.SetInput ring_matrix, ray_matrix,  success ...\n");

  std::uint32_t token_id_image{};
  std::uint32_t token_id_grid_sample{};

  std::uint32_t start_time = CV_FlexidagTimestamp();
  memset(grid_buffer_, 0, kGridSampleResSize);

  std::uint32_t Rval = grid_sample.RunNonBlock(GridSampleCallback, grid_buffer_,
                                               token_id_grid_sample);
  Print(".image_out_buf_.cat2491 = %p\n", image_out_buf_.cat2491);
  Rval = image_encoder.RunNonBlock(ImageEncoderCallback, &image_out_buf_,
                                   token_id_image);
  if (Rval != 0U) {
    Print("error: %d", Rval);
  }
  Print("RunNonBlock,  ... ...\n");

  Rval |= grid_sample.WaitRunFinish(token_id_grid_sample);
  if (Rval != 0U) {
    Print("WaitRunFinish error: %d", Rval);
  }
  Rval |= image_encoder.WaitRunFinish(token_id_image);

  if (Rval != 0U) {
    Print("WaitRunFinish error: %d", Rval);
  }
  Print("WaitRunFinish,  success ...\n");

  std::uint32_t end_time = CV_FlexidagTimestamp();
  //
  success = bev_net.SetInput(image_out_buf_.cat2491, kCat2491,
                             BevNet::InputType::CAT2491);
  success = bev_net.SetInput(image_out_buf_.cat2492, kCat2492,
                             BevNet::InputType::CAT2492);
  success = bev_net.SetInput(grid_buffer_, kGridSampleResSize,
                             BevNet::InputType::GRID_SAMPLE);
  if (!success) {
    Print("Error encounted when load input file bev net! Quit...\n");
    return 1;
  }
  Print("bev_net.SetInput.grid_buffer_-  success ...\n");
  bev_net.alloc_IO_buffer();
  auto end_time_2 = CV_FlexidagTimestamp();
  Rval = bev_net.RunBlock();
  auto end_time_3 = CV_FlexidagTimestamp();

  auto const& out_data = bev_net.GetOutBuf();
  out_data.num_of_buf
  /*
   "OUT0": {
        "box_0_cls,box_1_height,box_2_height,box_0_reg,box_0_vel,box_1_reg,box_3_cls,box_1_vel,box_1_cls,box_2_vel,box_3_reg,box_3_rot,box_3_vel":
   { "file": "combine_1_13_bub.out"
    */

  Print("**************************************\n");
  Print("Time of >> PART 1 << : %d us.\n", end_time - start_time);  // 5728  us
  Print("Time of transfering data: %d us.\n",
        end_time_2 - end_time);  // 576   us
  Print("Time of >> PART 2 << : %d us.\n",
        end_time_3 - end_time_2);  // 41408 us
  Print("Total inference time : %d us.\n",
        end_time_3 - start_time);  // 47712 us
  Print("**************************************\n");
}

void BevFusion::GridSampleCallback(void *vpHandle, uint32_t flexidag_output_num,
                                   flexidag_memblk_t *pblk_Output,
                                   void *vpParameter) {
  Print("begin ...flexidag_output_num = %d\n", flexidag_output_num);
  char *ptr = (char *)vpParameter;
  printf("GridSampleCallback ptr = %p\n", ptr);
  uint32_t Rval = AmbaCV_UtilityCmaMemInvalid(pblk_Output);
  if (Rval != 0U) {
    Print("Error encounted when GridSampleCallback CmaMemInvalid\n");
  }
  Print("pblk_Output->pBuffer = %p\n", pblk_Output->pBuffer);
  if (flexidag_output_num == 0) {
    memcpy(ptr, pblk_Output->pBuffer, kGridSampleResSize);
  }
  Print("end ... flexidag_output_num = %d\n", flexidag_output_num);
}

void BevFusion::ImageEncoderCallback(void *vpHandle,
                                     uint32_t flexidag_output_num,
                                     flexidag_memblk_t *pblk_Output,
                                     void *vpParameter) {
  Print("begin ... flexidag_output_num = %d\n", flexidag_output_num);
  ImageOutbuffer *ptr = (ImageOutbuffer *)vpParameter;
  printf(".cat2491 = %p\n", ptr->cat2491);
  uint32_t Rval = AmbaCV_UtilityCmaMemInvalid(pblk_Output);
  if (Rval != 0U) {
    Print("Error encounted when ImageEncoderCallback CmaMemInvalid\n");
  }
  Print("pblk_Output->pBuffer = %p\n", pblk_Output->pBuffer);

  if (flexidag_output_num == 0) {  // cat2491
    Print("memcpy cat2491 begin \n");
    memcpy(ptr->cat2491, pblk_Output->pBuffer, kCat2491);
    Print("memcpy cat2491 success\n");
  } else if (flexidag_output_num == 1) {  // 2492
    memcpy(ptr->cat2492, pblk_Output->pBuffer, kCat2492);
  }
}

uint32_t BevFusion::Test(Config const &config) noexcept {
  uint32_t Rval = 0;
  for (std::uint8_t i = 0; i < 6; ++i) {
    Rval |= image_encoder.SetImageFromFile(config.input_image_path, i);
  }

  Rval |= grid_sample.SetInput(config.pre_input_path,
                               GridSample::InputType::PRE_INPUT);
  Rval |= grid_sample.SetInput(config.align_matrix_path,
                               GridSample::InputType::ALIGN_MATRIX);

  bool success{true};
  success =
      bev_net.SetInput(config.ring_matrix_path, BevNet::InputType::RING_MATRIX);
  success =
      bev_net.SetInput(config.ray_matrix_path, BevNet::InputType::RAY_MATRIX);
  success =
      bev_net.SetInput(config.lidar_input_path, BevNet::InputType::LIDAR_INPUT);

  if (!success) {
    Print("Error encounted when load input file bev net! Quit...\n");
    return 1;
  }

  std::uint32_t token_id_image{};
  std::uint32_t token_id_grid_sample{};

  std::uint32_t start_time = CV_FlexidagTimestamp();
  memset(grid_buffer_, 0, kGridSampleResSize);

  Rval =
      grid_sample.RunNonBlock(GridSampleCallback, grid_buffer_, token_id_image);
  Rval = image_encoder.RunNonBlock(ImageEncoderCallback, &image_out_buf_,
                                   token_id_grid_sample);
  if (Rval != 0U) {
    Print("error: %d\n", Rval);
  }

  Rval |= grid_sample.WaitRunFinish(token_id_grid_sample);
  Print("token_id_grid_sample = %d\n", token_id_grid_sample);
  Rval |= image_encoder.WaitRunFinish(token_id_image);

  if (Rval != 0U) {
    Print("error: %d", Rval);
  }
  std::uint32_t end_time = CV_FlexidagTimestamp();
  //
  success = bev_net.SetInput(image_out_buf_.cat2491, kCat2491,
                             BevNet::InputType::CAT2491);
  success = bev_net.SetInput(image_out_buf_.cat2492, kCat2492,
                             BevNet::InputType::CAT2492);
  success = bev_net.SetInput(grid_buffer_, kGridSampleResSize,
                             BevNet::InputType::GRID_SAMPLE);
  if (!success) {
    Print("Error encounted when load input file bev net! Quit...\n");
    return 1;
  }

  auto end_time_2 = CV_FlexidagTimestamp();
  Rval = bev_net.RunBlock();
  auto end_time_3 = CV_FlexidagTimestamp();

  Print("**************************************\n");
  Print("Time of >> PART 1 << : %d us.\n", end_time - start_time);  // 5728  us
  Print("Time of transfering data: %d us.\n",
        end_time_2 - end_time);  // 576   us
  Print("Time of >> PART 2 << : %d us.\n",
        end_time_3 - end_time_2);  // 41408 us
  Print("Total inference time : %d us.\n",
        end_time_3 - start_time);  // 47712 us
  Print("**************************************\n");
}
